from io import *
import os

class DataBase(object):
    def __init__(self, dir_name):
        self.dir_name = dir_name

    def get_count(self, file):
        db = open(os.path.join(self.dir_name, file))
        ans = 0
        base = []
        line = db.readline()
        while line != '==========\n':
            base.append(line)
            if line.split('\t')[0] == 'count_of_lines':
                ans = int(line.split('\t')[1])
            line = db.readline()
        base.append(line)
        return ans, base, db
    
    def skip(self, file):
        db = open(os.path.join(self.dir_name, file))
        base = []
        line = db.readline()
        while line != '==========\n':
            base.append(line)
            line = db.readline()
        base.append(line)
        return base, db

    def add(self, file, task_to_add):
        count_of_lines, base, db = self.get_count(file)
        count_of_lines += 1
        for i in range(len(base)):
            base[i] = list(base[i].split('\t'))
            if base[i][0] == 'count_of_lines':
                base[i][1] = str(count_of_lines) + '\n'
            base[i] = '\t'.join(base[i])
        is_in_db = False
        task = db.readline().rstrip()
        #while task != '':
            #base.append(task + '\n')
            #task = list(task.split('\t'))
            #if len(task) >= 2 and task[1] == task_to_add:
                #is_in_db = True
            #task = db.readline().rstrip()
        db.close()
        db = open(os.path.join(self.dir_name, file), 'w')
        for task in base:
            db.write(task)
        if not is_in_db:
            db.write(unicode(count_of_lines) + '\t' + unicode(task_to_add) + '\n')
        db.close()

    def delete(self, file, pk):
        base, db = self.skip(file)
        task = db.readline().rstrip()
        while task != '':
            if int(task.split('\t')[0]) == pk:
                base.append('')
            else:
                base.append(task + '\n')
            task = db.readline().rstrip()
        db.close()
        db = open(os.path.join(self.dir_name, file), 'w')
        for i in range(len(base)):
            db.write(unicode(base[i]))
        db.close()

    def get(self, file, pk):
        base, db = self.skip(file)
        ans = []
        for task in db:
            if task.split('\t')[0] == pk:
                ans.append(task.split('\t')[1])
        if ans != []:
            return ans[0]
        else:
            return None

    def extract_all(self, file):
        base, db = self.skip(file)
        ans = [task.split('\t')[1].rstrip() for task in db]
        db.close()
        return ans
